package com.mm.base.filter.xss2;

import com.mm.base.exception.MuMuException;
import lombok.extern.slf4j.Slf4j;

import javax.servlet.ReadListener;
import javax.servlet.ServletInputStream;
import javax.servlet.ServletOutputStream;
import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletRequestWrapper;
import javax.servlet.http.HttpServletResponse;
import java.io.BufferedReader;
import java.io.ByteArrayInputStream;
import java.io.IOException;
import java.io.InputStreamReader;
import java.nio.charset.Charset;
import java.util.regex.Matcher;
import java.util.regex.Pattern;

import static java.util.regex.Pattern.*;

/**
 * @Description Xss请求过滤类
 * @Date 2024/4/11 下午1:44
 * @Author yanglin
 **/
@Slf4j
public class XssHttpServletRequestWrapper extends HttpServletRequestWrapper {

    public XssHttpServletRequestWrapper() {
        super(null);
    }

    /**
     * Constructs a request object wrapping the given request.
     *
     * @param request The request to wrap
     * @throws IllegalArgumentException if the request is null
     */
    public XssHttpServletRequestWrapper(HttpServletRequest request) {
        super(request);
    }

    /**
     * 过滤springmvc中的 @RequestParam 注解中的参数
     *
     * @param s a <code>String</code> containing the name of the parameter
     *          whose value is requested
     * @return
     */
    @Override
    public String[] getParameterValues(String s) {
        String[] paramStr = super.getParameterValues(s);
        if (paramStr == null) {
            return null;
        }
        int i = paramStr.length;
        String[] param = new String[i];
        for (int j = 0; j < i; j++) {
            param[j] = cleanXSS(cleanSQLInject(paramStr[j]));
        }
        log.info("XssHttpServletRequestWrapper净化后的请求为：{}", param);
        return param;
    }

    /**
     * 过滤request.getParameter的参数
     *
     * @param s a <code>String</code> specifying the name of the parameter
     * @return
     */
    @Override
    public String getParameter(String s) {
        String paramStr = super.getParameter(s);
        if (paramStr == null) {
            return null;
        } else {
            String param = cleanXSS(cleanSQLInject(paramStr));
            log.info("XssHttpServletRequestWraper净化后的请求为：{}", param);
            return param;
        }
    }

    /**
     * 过滤请求体 json 格式的
     *
     * @return
     * @throws IOException
     */
    @Override
    public ServletInputStream getInputStream() throws IOException {
        final ByteArrayInputStream byteArrayInputStream = new ByteArrayInputStream(inputHandlers(super.getInputStream()).getBytes());
        return new ServletInputStream() {
            @Override
            public int read() throws IOException {
                return byteArrayInputStream.read();
            }

            @Override
            public boolean isFinished() {
                return false;
            }

            @Override
            public boolean isReady() {
                return false;
            }

            @Override
            public void setReadListener(ReadListener readListener) {
            }
        };
    }


    public String inputHandlers(ServletInputStream servletInputStream) {
        StringBuilder sb = new StringBuilder();
        BufferedReader reader = null;
        try {
            reader = new BufferedReader(new InputStreamReader(servletInputStream, Charset.forName("UTF-8")));
            String line = "";
            while ((line = reader.readLine()) != null) {
                sb.append(line);
            }
        } catch (IOException e) {
            e.printStackTrace();
        } finally {
            if (servletInputStream != null) {
                try {
                    servletInputStream.close();
                } catch (IOException e) {
                    e.printStackTrace();
                }
            }
            if (reader != null) {
                try {
                    reader.close();
                } catch (IOException e) {
                    e.printStackTrace();
                }
            }
        }
        return cleanXSS(sb.toString());
    }

    /**
     * 处理XSS参数
     *
     * @param xssStr
     * @return
     */
    public String cleanXSS(String xssStr) {
        String temp = xssStr;
        xssStr = xssStr.replaceAll("<", "＜").replaceAll(">", "＞");
        xssStr = xssStr.replaceAll("\\(", "（").replaceAll("\\)", "）");
        xssStr = xssStr.replaceAll("'", "＇");
        xssStr = xssStr.replaceAll(";", "；");
        /**-----------------------start--------------------------*/
        xssStr = xssStr.replaceAll("<", "& lt;").replaceAll(">", "& gt;");
        xssStr = xssStr.replaceAll("\\(", "& #40;").replaceAll("\\)", "& #41");
        xssStr = xssStr.replaceAll("eval\\((.*)\\)", "");
        xssStr = xssStr.replaceAll("[\\\"\\\'][\\s]*javascript:(.*)[\\\"\\\']", "\"\"");
        xssStr = xssStr.replaceAll("script", "");
        xssStr = xssStr.replaceAll("link", "");
        xssStr = xssStr.replaceAll("frame", "");
        /**-----------------------end--------------------------*/
        Pattern pattern = compile("(eval\\((.*)\\)|script)",
                CASE_INSENSITIVE);
        Matcher matcher = pattern.matcher(xssStr);
        xssStr = matcher.replaceAll("");

        pattern = compile("[\\\"\\'][\\s]*javascript:(.*)[\\\"\\']",
                CASE_INSENSITIVE);
        matcher = pattern.matcher(xssStr);
        xssStr = matcher.replaceAll("\"\"");
        // 增加脚本
        xssStr = xssStr.replaceAll("script", "").replaceAll(";", "")
                /*.replaceAll("\"", "").replaceAll("@", "")*/
                .replaceAll("0x0d", "").replaceAll("0x0a", "");

        if (!temp.equals(xssStr)) {
            log.error("xss攻击检查：参数含有非法攻击字符，已禁止继续访问！！");
            log.error("原始输入信息-->" + temp);

            throw new MuMuException("xss攻击检查：参数含有非法攻击字符，已禁止继续访问！！");
        }
        return xssStr;
    }

    /**
     * 输出
     *
     * @param response
     * @param msg
     * @throws IOException
     */
    public void outputMsgByOutputStream(HttpServletResponse response, String msg) throws IOException {
        // 获取输出流
        ServletOutputStream outputStream = response.getOutputStream();
        // 通过设置响应头控制浏览器以UTF-8的编码显示数据，如果不加这句话，那么浏览器显示的将是乱码
        response.setHeader("content-type", "text/html;charset=UTF-8");
        // 将字符转换成字节数组，指定以UTF-8 编码进行转换
        byte[] dataByteArr = msg.getBytes("UTF-8");
        // 使用OutputStream流向客户端输出字节数组
        outputStream.write(dataByteArr);
    }

    /**
     * 需要增加通配，过滤大小写组合
     *
     * @param xssStr
     * @return
     */
    public String cleanSQLInject(String xssStr) {
        String lowXssStr = xssStr.toLowerCase();
        String temp = xssStr;
        String lowXssStrAfter = lowXssStr.replaceAll("insert", "forbidI")
                .replaceAll("select", "forbidS")
                .replaceAll("update", "forbidU")
                .replaceAll("delete", "forbidD").replaceAll("and", "forbidA")
                .replaceAll("or", "forbidO");

        if (!lowXssStrAfter.equals(lowXssStr)) {
            log.error("sql注入检查：输入信息存在SQL攻击！");
            log.error("原始输入信息-->" + temp);
            log.error("处理后信息-->" + lowXssStr);
            throw new MuMuException("sql注入检查：参数含有非法攻击字符，已禁止继续访问！！");
        }
        return xssStr;
    }
}
